package com.knxhd.config;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import com.alibaba.fastjson.JSONObject;
import com.knxhd.enums.Role;
import com.knxhd.pojo.LlmRequest;
import com.knxhd.pojo.LlmResponse;
import com.knxhd.pojo.LlmResponse.Choice;
import com.knxhd.pojo.LlmResponse.MessageDetail;
import com.knxhd.pojo.LlmResponse.Usage;

import cn.hutool.core.collection.ListUtil;
import cn.hutool.core.util.RandomUtil;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.Content;
import dev.langchain4j.data.message.ContentType;
import dev.langchain4j.data.message.TextContent;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.listener.ChatModelErrorContext;
import dev.langchain4j.model.chat.listener.ChatModelListener;
import dev.langchain4j.model.chat.listener.ChatModelRequestContext;
import dev.langchain4j.model.chat.listener.ChatModelResponseContext;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.request.ChatRequestParameters;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.ChatResponseMetadata;
import dev.langchain4j.model.openai.OpenAiChatRequestParameters;
import dev.langchain4j.model.openai.OpenAiChatResponseMetadata;
import dev.langchain4j.model.openai.OpenAiTokenUsage;
import dev.langchain4j.model.output.TokenUsage;
import lombok.extern.slf4j.Slf4j;

/**
 * Author: tianluhua tianlh@qdcares.cn
 * Date: 2025-07-13 10:11:39
 * Description: langChain4j配置类
 */
@Configuration
@Slf4j
public class Langchain4jConfig {

  @Bean
  public ChatModelListener chatModelListener() {
    return new ChatModelListener() {

      @Override
      public void onError(ChatModelErrorContext errorContext) {
        log.error("大模型调用异常");
        Throwable error = errorContext.error();
        error.printStackTrace();

        ChatRequest chatRequest = errorContext.chatRequest();
        System.out.println(chatRequest);

        System.out.println(errorContext.modelProvider());

        Map<Object, Object> attributes = errorContext.attributes();
        System.out.println(attributes.get("my-attribute"));
      }

      @Override
      public void onRequest(ChatModelRequestContext requestContext) {
        String sessionId = RandomUtil.randomString(64);
        requestContext.attributes().put("sessionId", sessionId);
        ChatRequest chatRequest = requestContext.chatRequest();

        List<ChatMessage> messages = chatRequest.messages();
        List<MessageDetail> messageDetails = new ArrayList<>();

        LlmRequest llmRequest = new LlmRequest();
        for (ChatMessage chatMessage : messages) {
          UserMessage userMessage = (UserMessage) chatMessage;
          List<Content> contents = userMessage.contents();
          for (Content content : contents) {
            if (ContentType.TEXT.equals(content.type())) {
              TextContent textContent = (TextContent) content;
              messageDetails.add(new MessageDetail(Role.user, textContent.text()));
            }
          }
        }

        ChatRequestParameters parameters = chatRequest
            .parameters();
        llmRequest.setMessages(messageDetails);
        llmRequest.setModelName(parameters.modelName());
        llmRequest.setTemperature(parameters.temperature());
        llmRequest.setTopP(parameters.topP());
        llmRequest.setTopK(parameters.topK());
        llmRequest.setFrequencyPenalty(parameters.frequencyPenalty());
        llmRequest.setPresencePenalty(parameters.presencePenalty());
        llmRequest.setMaxOutputTokens(parameters.maxOutputTokens());
        llmRequest.setStopSequences(parameters.stopSequences());

        System.out.println(parameters.modelName());
        System.out.println(parameters.temperature());
        System.out.println(parameters.topP());
        System.out.println(parameters.topK());
        System.out.println(parameters.frequencyPenalty());
        System.out.println(parameters.presencePenalty());
        System.out.println(parameters.maxOutputTokens());
        System.out.println(parameters.stopSequences());
        System.out.println(parameters.toolSpecifications());
        System.out.println(parameters.toolChoice());
        System.out.println(parameters.responseFormat());

        if (parameters instanceof OpenAiChatRequestParameters openAiParameters) {
          System.out.println(openAiParameters.maxCompletionTokens());
          System.out.println(openAiParameters.logitBias());
          System.out.println(openAiParameters.parallelToolCalls());
          System.out.println(openAiParameters.seed());
          System.out.println(openAiParameters.user());
          System.out.println(openAiParameters.store());
          System.out.println(openAiParameters.metadata());
          System.out.println(openAiParameters.serviceTier());
          System.out.println(openAiParameters.reasoningEffort());
        }

        System.out.println(requestContext.modelProvider());

        Map<Object, Object> attributes = requestContext.attributes();
        attributes.put("my-attribute", "my-value");

        log.info("大模型调用参数: sessionId: {}, 请求参数:{}", sessionId, JSONObject.toJSONString(llmRequest));
      }

      @Override
      public void onResponse(ChatModelResponseContext responseContext) {
        ChatResponse chatResponse = responseContext.chatResponse();

        AiMessage aiMessage = chatResponse.aiMessage();
        System.out.println(aiMessage);

        ChatResponseMetadata metadata = chatResponse.metadata();
        System.out.println(metadata.id());
        System.out.println(metadata.modelName());
        System.out.println(metadata.finishReason());

        if (metadata instanceof OpenAiChatResponseMetadata openAiMetadata) {
          System.out.println(openAiMetadata.created());
          System.out.println(openAiMetadata.serviceTier());
          System.out.println(openAiMetadata.systemFingerprint());
        }

        TokenUsage tokenUsage = metadata.tokenUsage();
        System.out.println(tokenUsage.inputTokenCount());
        System.out.println(tokenUsage.outputTokenCount());
        System.out.println(tokenUsage.totalTokenCount());
        if (tokenUsage instanceof OpenAiTokenUsage openAiTokenUsage) {
          if (openAiTokenUsage.inputTokensDetails() != null) {
            System.out.println(openAiTokenUsage.inputTokensDetails().cachedTokens());
          }
          if (openAiTokenUsage.outputTokensDetails() != null) {
            System.out.println(openAiTokenUsage.outputTokensDetails().reasoningTokens());
          }
        }

        ChatRequest chatRequest = responseContext.chatRequest();
        System.out.println(chatRequest);

        System.out.println(responseContext.modelProvider());

        Map<Object, Object> attributes = responseContext.attributes();
        System.out.println(attributes.get("my-attribute"));

        String sessionId = responseContext.attributes().get("sessionId") + "";
        LlmResponse llmResponse = LlmResponse.builder()
            .choices(
                ListUtil.of(new Choice(new MessageDetail(Role.getRole(aiMessage.type().name()), aiMessage.text()))))
            .usage(new Usage(tokenUsage.outputTokenCount(), tokenUsage.inputTokenCount(), tokenUsage.totalTokenCount()))
            .build();

        log.info("大模型返回结果: sessionId: {}, 请求参数:{}", sessionId, JSONObject.toJSONString(llmResponse));

      }

    };
  }

}
